package com.starhub.application.model.biz;

import com.starhub.integration.sdk.baidu.model.BaiduEmbeddings;
import com.starhub.integration.sdk.custom.model.CustomEmbeddings;

import dev.langchain4j.model.embedding.EmbeddingModel;

import com.starhub.application.interfaces.model.ModelConfigService;
import com.starhub.application.model.enums.VectorModelTypeEnum;
import com.starhub.common.bean.model.ModelConfig;
import com.starhub.common.security.util.SecurityUtils;
import com.starhub.common.sys.user.LoginUser;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

/**
 * 向量模型工厂类
 */
@Component
public class VectorModelFactory {

    @Autowired
    private ModelConfigService modelConfigService;

    /**
     * 获取向量模型实例
     *
     * @param mark 向量模型标识
     * @return 向量模型实例
     */
    public EmbeddingModel getVectorModel(String mark) {
        // 从数据库获取模型配置
        LoginUser user = SecurityUtils.getCurrentUser();
        ModelConfig config = modelConfigService.getModel(user, mark);
        if (config == null) {
            throw new IllegalArgumentException("未找到向量模型配置: " + mark);
        }
        VectorModelTypeEnum modelType = VectorModelTypeEnum.getByCode(config.getChatType());
        // 根据模型类型创建对应的实例
        switch (modelType) {
            case EmbeddingV1:
                return BaiduEmbeddings.builder()
                        .mark(config.getMark())
                        .apiKey(config.getApiKey())
                        .secretKey(config.getApiSecret())
                        .build();
            case Custom:
                //自定义的向量模型apiKey存储模型标识
                return CustomEmbeddings.builder()
                        .mark(config.getMark())
                        .apiUrl(config.getApiUrl())
                        .modelName(config.getApiKey())
                        .build();
            default:
                throw new IllegalArgumentException("不支持的向量模型类型: " + modelType);
        }
    }
}